import json
import shutil
import base64
import random
import string
import io
import requests
import datetime
import os
import logging
import uuid

from flask_restplus import Namespace, Resource, reqparse
from flask_login import login_required, current_user
from werkzeug.datastructures import FileStorage
from flask import send_file
from mongoengine.errors import NotUniqueError
from PIL import Image


from ..util import query_util, coco_util
#from ..util.watermark_util import embed_watermark_str, extract_watermark_str
from database import (
    ImageModel,
    DatasetModel,
    AnnotationModel,
    CategoryModel
)


logger = logging.getLogger('gunicorn.error')

api = Namespace('image', description='Image related operations')

image_all = reqparse.RequestParser()
image_all.add_argument('fields', required=False, type=str)
image_all.add_argument('page', default=1, type=int)
image_all.add_argument('per_page', default=50, type=int, required=False)

image_upload = reqparse.RequestParser()
image_upload.add_argument('image', location='files',
                          type=FileStorage, required=True,
                          help='PNG or JPG file')
image_upload.add_argument('dataset_id', required=True, type=int,
                          help='Id of dataset to insert image into')

image_download = reqparse.RequestParser()
image_download.add_argument('asAttachment', type=bool, default=False)
image_download.add_argument('thumbnail', type=bool, default=False)
image_download.add_argument('width', type=int)
image_download.add_argument('height', type=int)

image_enhance = reqparse.RequestParser()
# image_enhance.add_argument('coords_type',type=str)
image_enhance.add_argument('image_id', type=int, required=True, location="json")
image_enhance.add_argument('point_coords', type=list, location="json")
image_enhance.add_argument("point_labels", type=list, location="json")
image_enhance.add_argument("text_prompt", type=str, location="json")

copy_annotations = reqparse.RequestParser()
copy_annotations.add_argument('category_ids', location='json', type=list,
                              required=False, default=None, help='Categories to copy')

image_to_inpaint = reqparse.RequestParser()
image_to_inpaint.add_argument('image_id', type=int, required=True)
# image_enhance.add_argument('coords_type',type=str)
image_to_inpaint.add_argument('text_prompt', type=str, default=None, required=True)
image_to_inpaint.add_argument("point_coords", type=list, required=True)
image_to_inpaint.add_argument("point_labels", type=int)
# image_to_inpaint.add_argument('image', location='files',
#                           type=FileStorage, required=True,
#                           help='PNG or JPG file')
# 前端要自动获取的数据集id
# image_to_inpaint.add_argument('dataset_id', required=True, type=int,
#                               help='Id of dataset to insert image into')


image_deleted_by_id = reqparse.RequestParser()
image_deleted_by_id.add_argument('image_id', type=int, required=True)

image_save_or_delete = reqparse.RequestParser()
image_save_or_delete.add_argument('image_id', type=int, required=True)
image_save_or_delete.add_argument('image_origin_id', type=int, required=True)
image_save_or_delete.add_argument('decision', type=bool, required=True)


image_rename_args = reqparse.RequestParser()
image_rename_args.add_argument('image_id', type=int, required=True, help='Image ID is required')
image_rename_args.add_argument('new_filename', type=str, required=True, help='New filename is required')


image_remove_by_segmentatino = reqparse.RequestParser()
image_remove_by_segmentatino.add_argument('image_id', type=int, required=True)
image_remove_by_segmentatino.add_argument("annotation_id", type=int)


image_fill_by_segmentatino = reqparse.RequestParser()
image_fill_by_segmentatino.add_argument('image_id', type=int, required=True)
image_fill_by_segmentatino.add_argument('annotation_id', type=int)
image_fill_by_segmentatino.add_argument('text_prompt', type=str)

# embed_watermark_str(input_img_path, output_img_path)
embed_watermark = reqparse.RequestParser()
embed_watermark.add_argument('input_img_id', type=int, required=True)
# embed_watermark.add_argument('output_img_path', type=str, required=True)

# extract_watermark_str(input_img_with_watermark_path)
extract_watermark = reqparse.RequestParser()
extract_watermark.add_argument('input_img_with_watermark_id', type=int, required=True)

pix_to_pix = reqparse.RequestParser()
pix_to_pix.add_argument('image_id', type=int, required=True)
pix_to_pix.add_argument('text_prompt', type=str, default=None, required=True)


@api.route('/')
class Images(Resource):

    @api.expect(image_all)
    @login_required
    def get(self):
        """ Returns all images """
        args = image_all.parse_args()
        per_page = args['per_page']
        page = args['page'] - 1
        fields = args.get('fields', '')

        images = current_user.images.filter(deleted=False)
        total = images.count()
        pages = int(total / per_page) + 1

        images = images.skip(page * per_page).limit(per_page)
        if fields:
            images = images.only(*fields.split(','))

        return {
            "total": total,
            "pages": pages,
            "page": page,
            "fields": fields,
            "per_page": per_page,
            "images": query_util.fix_ids(images.all())
        }

    @api.expect(image_upload)
    @login_required
    def post(self):
        """ Creates an image """
        args = image_upload.parse_args()
        image = args['image']

        dataset_id = args['dataset_id']
        try:
            dataset = DatasetModel.objects.get(id=dataset_id)
        except:
            return {'message': 'dataset does not exist'}, 400
        directory = dataset.directory
        path = os.path.join(directory, image.filename)
        if os.path.exists(path):
            return {'message': 'file already exists'}, 400
        pil_image = Image.open(io.BytesIO(image.read()))
        pil_image.save(path)
        image.close()
        pil_image.close()
        try:
            db_image = ImageModel.create_from_path(path, dataset_id).save()
        except NotUniqueError:
            db_image = ImageModel.objects.get(path=path)
        return db_image.id


@api.route('/<int:image_id>')
class ImageId(Resource):

    @api.expect(image_download)
    @login_required
    def get(self, image_id):
        """ Returns category by ID """
        args = image_download.parse_args()
        as_attachment = args.get('asAttachment')
        thumbnail = args.get('thumbnail')

        image = current_user.images.filter(id=image_id, deleted=False).first()

        if image is None:
            return {'success': False}, 400

        width = args.get('width')
        height = args.get('height')

        if not width:
            width = image.width
        if not height:
            height = image.height

        pil_image = image.open_thumbnail() if thumbnail else Image.open(image.path)

        pil_image.thumbnail((width, height), Image.ANTIALIAS)
        image_io = io.BytesIO()
        pil_image = pil_image.convert("RGB")
        pil_image.save(image_io, "JPEG", quality=90)
        image_io.seek(0)

        return send_file(image_io, attachment_filename=image.file_name, as_attachment=as_attachment)

    @login_required
    def delete(self, image_id):
        """ Deletes an image by ID """
        image = current_user.images.filter(id=image_id, deleted=False).first()
        if image is None:
            return {"message": "Invalid image id"}, 400

        if not current_user.can_delete(image):
            return {"message": "You do not have permission to download the image"}, 403

        image.update(set__deleted=True, set__deleted_date=datetime.datetime.now())
        return {"success": True}


@api.route('/copy/<int:from_id>/<int:to_id>/annotations')
class ImageCopyAnnotations(Resource):

    @api.expect(copy_annotations)
    @login_required
    def post(self, from_id, to_id):
        args = copy_annotations.parse_args()
        category_ids = args.get('category_ids')

        image_from = current_user.images.filter(id=from_id).first()
        image_to = current_user.images.filter(id=to_id).first()

        if image_from is None or image_to is None:
            return {'success': False, 'message': 'Invalid image ids'}, 400

        if image_from == image_to:
            return {'success': False, 'message': 'Cannot copy self'}, 400

        if image_from.width != image_to.width or image_from.height != image_to.height:
            return {'success': False, 'message': 'Image sizes do not match'}, 400

        if category_ids is None:
            category_ids = DatasetModel.objects(id=image_from.dataset_id).first().categories

        query = AnnotationModel.objects(
            image_id=image_from.id,
            category_id__in=category_ids,
            deleted=False
        )

        return {'annotations_created': image_to.copy_annotations(query)}


@api.route('/<int:image_id>/coco')
class ImageCoco(Resource):

    @login_required
    def get(self, image_id):
        """ Returns coco of image and annotations """
        image = current_user.images.filter(id=image_id).exclude('deleted_date').first()

        if image is None:
            return {"message": "Invalid image ID"}, 400

        if not current_user.can_download(image):
            return {"message": "You do not have permission to download the images's annotations"}, 403

        return coco_util.get_image_coco(image_id)


def load_image(image_path):
    try:
        img = Image.open(image_path)
        return img
    except IOError:
        print("无法打开或识别图像文件：", image_path)
        return None


def random_string(length):
    """生成指定长度的随机字符串"""
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for _ in range(length))

def rename_file(file_path):
    """将指定图片文件名除了后缀之外的部分改成随机字符串"""
    if os.path.isfile(file_path) and file_path.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp')):
        file_directory, file_name = os.path.split(file_path)
        file_name_without_extension, file_extension = os.path.splitext(file_name)
        new_file_name = random_string(8) + file_extension
        os.rename(file_path, os.path.join(file_directory, new_file_name))
        print(f"Renamed {file_name} to {new_file_name}")
    else:
        print(f"Error: {file_path} is not a valid image file.")


@api.route('/replace_anything')
class ImageInpaint(Resource):

    @api.expect(image_to_inpaint)
    def post(self):
        args = image_to_inpaint.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        image_path = image.path  # 将路径替换为实际图像文件的路径
        image = load_image(image_path)
        if not image:
            return {"message": "Invalid image ID"}, 400
        logging.info("1" * 10)
        remote_flask_url = 'http://125.220.157.228:29990/inpaint/upload'
        response_image_saved_path_and_name = requests.post(remote_flask_url, files=image)
        logging.info("2" * 10)
        response_path_and_name = json.loads(response_image_saved_path_and_name.text)
        file_path = response_path_and_name['message']
        data = {'input_img': file_path,
                'coords_type': "key_in",
                'point_coords': args.getlist('point_coords'),
                'point_labels': args['point_labels'],
                'text_prompt': args['text_prompt'],
                'dilate_kernel_size': 10,  # fixme
                'output_dir': './results'
                }
        remote_flask_inpaint_replace_anything_url = 'http://125.220.157.228:29990/inpaint/replace_thing'
        logging.info("remote_flask_inpaint_replace_anything_url" + remote_flask_inpaint_replace_anything_url)
        response_replace_thing = requests.post(remote_flask_inpaint_replace_anything_url, data=data)
        response_replace_thing = json.loads(response_replace_thing.text)
        response_replace_thing_image_bytes = response_replace_thing['image']
        dataset_id = 1
        try:
            dataset = DatasetModel.objects.get(id=dataset_id)
        except:
            return {'message': 'dataset does not exist'}, 400
        directory = dataset.directory
        path = os.path.join(directory, response_replace_thing_image_bytes.filename)
        if os.path.exists(path):
            return {'message': 'file already exists'}, 400
        pil_image = Image.open(io.BytesIO(response_replace_thing_image_bytes.read()))
        pil_image.save(path)
        image.close()
        pil_image.close()
        try:
            db_image = ImageModel.create_from_path(path, dataset_id).save()
        except NotUniqueError:
            db_image = ImageModel.objects.get(path=path)
        return db_image.id


def upload_image_to_remote_flask(local_image_path, remote_flask_url):
    try:
        # 打开本地图片文件
        with open(local_image_path, 'rb') as f:
            # 准备发送的数据
            files = {'file': f}
            # 发送POST请求到远程Flask应用
            response = requests.post(remote_flask_url, files=files)
            # 返回响应内容
            return response.text
    except Exception as e:
        return str(e)


@api.route('/enhance')
class ImageEnhance(Resource):

    @api.expect(image_enhance)
    def post(self):
        args = image_enhance.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        if not image:
            return {"message": "Invalid image ID"}, 400
        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory
        enhance_directory = dataset_path + r'enhance/'
        if not os.path.exists(enhance_directory):
            os.makedirs(enhance_directory)
        enhance_filepath = os.path.join(enhance_directory, image.file_name)

        image_path = image.path  # 将路径设置为为实际图像文件的路径
        image = load_image(image_path)
        # files = {'file': image}
        logger.info(image)
        logger.info("6" * 10)

        remote_flask_url = 'http://125.220.157.228:29990/inpaint/upload'
        response_image_saved_path_and_name = upload_image_to_remote_flask(image_path, remote_flask_url)
        # response_image_saved_path_and_name = requests.post(remote_flask_url, files=files)
        logger.info("7" * 10)
        logger.info("response_image_saved_path_and_name" + response_image_saved_path_and_name)
        response_data = json.loads(response_image_saved_path_and_name)
        file_path = response_data["message"]
        logger.info("file_path" + "---" + file_path)
        logger.info("8" * 10)
        data = {'input_img': file_path,
                'coords_type': "key_in",
                'point_coords': args['point_coords'],
                'point_labels': args['point_labels'],
                'text_prompt': args['text_prompt'],
                'dilate_kernel_size': 10,  # fixme
                'output_dir': './results'
                }
        logger.info("point_coords" + str(data['point_coords']))

        logger.info("9" * 10)

        remote_flask_inpaint_replace_anything_url = 'http://125.220.157.228:29990/inpaint/replace_thing'
        response_replace_thing = requests.post(remote_flask_inpaint_replace_anything_url, json=data)
        logger.info("10" * 10)
        logger.info("response_replace_thing: type: %s", type(response_replace_thing))
        logger.info("11" * 10)
        response_replace_thing = response_replace_thing.json()
        response_replace_thing_image_bytes = response_replace_thing['image']
        # logger.info("response_replace_thing_image_bytes" + response_replace_thing_image_bytes)
        logger.info("12" * 10)
        image_data = base64.b64decode(response_replace_thing_image_bytes)
        responsed_image = Image.open(io.BytesIO(image_data))
        logger.info("13" * 10)
        shutil.copyfile(image_path, enhance_filepath)
        responsed_image.save(enhance_filepath)
        image = ImageModel.create_from_path(enhance_filepath)
        image.enhanced = True
        image.save()
        return query_util.fix_ids(image)

@api.route('/fill')
class ImageFilled(Resource):

    @api.expect(image_enhance)
    def post(self):
        args = image_enhance.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        if not image:
            return {"message": "Invalid image ID"}, 400
        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory
        fill_directory = dataset_path + r'filled/'
        if not os.path.exists(fill_directory):
            os.makedirs(fill_directory)

        random_file_name = str(uuid.uuid4()) + ".jpg"
        fill_filepath = os.path.join(fill_directory, random_file_name)

        image_path = image.path  # 将路径设置为为实际图像文件的路径
        image = load_image(image_path)
        # files = {'file': image}
        logger.info(image)
        logger.info("6" * 10)

        remote_flask_url = 'http://125.220.157.228:29990/inpaint/upload'
        response_image_saved_path_and_name = upload_image_to_remote_flask(image_path, remote_flask_url)
        # response_image_saved_path_and_name = requests.post(remote_flask_url, files=files)
        logger.info("7" * 10)
        logger.info("response_image_saved_path_and_name" + response_image_saved_path_and_name)
        response_data = json.loads(response_image_saved_path_and_name)
        file_path = response_data["message"]
        logger.info("file_path" + "---" + file_path)
        logger.info("8" * 10)
        data = {'input_img': file_path,
                'coords_type': "key_in",
                'point_coords': args['point_coords'],
                'point_labels': args['point_labels'],
                'text_prompt': args['text_prompt'],
                'dilate_kernel_size': 10,  # fixme
                'output_dir': './results'
                }
        logger.info("point_coords" + str(data['point_coords']))

        remote_flask_fill_replace_anything_url = 'http://125.220.157.228:29990/inpaint/fill'
        response_fill_thing = requests.post(remote_flask_fill_replace_anything_url, json=data)
        logger.info("response_fill_thing: type: %s", type(response_fill_thing))
        response_fill_thing = response_fill_thing.json()
        response_fill_thing_image_bytes = response_fill_thing['image']
        image_data = base64.b64decode(response_fill_thing_image_bytes)
        responsed_image = Image.open(io.BytesIO(image_data))
        logger.info("13" * 10)
        shutil.copyfile(image_path, fill_filepath)
        responsed_image.save(fill_filepath)
        image = ImageModel.create_from_path(fill_filepath)
        image.enhanced = True
        image.deleted = True
        image.save()
        return query_util.fix_ids(image)


@api.route('/save_or_not')
class ImageEnhancedSaveOrDelete(Resource):
#
     @api.expect(image_save_or_delete)
     def post(self):
         args = image_save_or_delete.parse_args()
         decision = args['decision']
         image_id = args['image_id']
         image_origin_id = args['image_origin_id']
         if decision == True:
            image_origin = ImageModel.objects(id=image_origin_id).first()
            image_enhanced = ImageModel.objects(id=image_id).first()
            image_origin.deleted = False

            directory = os.path.dirname(image_origin.path)
            filename = os.path.basename(image_enhanced.path)
            new_path = os.path.join(directory, filename)
            # 复制图片文件
            shutil.copy(image_enhanced.path, new_path)

            # image_origin.path = new_path
            # image_enhanced.delete()
            image_enhanced.deleted=True
            image_origin.deleted=True
            # image_origin.update(path=new_path)
            image_enhanced.save()
            image_origin.save()

            new_image = ImageModel.create_from_path(new_path)
            new_image.update(
                annotated=image_origin.annotated,
                annotating=image_origin.annotating,
                file_name=image_origin.file_name,
                num_annotations=image_origin.num_annotations,
                regenerate_thumbnail=True,
                enhanced=True,
                metadata=image_origin.metadata
            )
            new_image.save()

            annotation_list = AnnotationModel.objects(image_id=image_id)
            for annotation in annotation_list:
                annotation.image_id = new_image.id
                annotation.save()

            return query_util.fix_ids(new_image)
         elif decision == False:
            return 'Image saved to another path successfully.'
         else:
            return 'Invalid decision.'


@api.route('/remove')
class ImageRemoved(Resource):

    @api.expect(image_enhance)
    def post(self):
        args = image_enhance.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        if not image:
            return {"message": "Invalid image ID"}, 400
        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory
        remove_directory = dataset_path + r'removed/'
        if not os.path.exists(remove_directory):
            os.makedirs(remove_directory)
        random_file_name = str(uuid.uuid4()) + ".jpg"
        remove_filepath = os.path.join(remove_directory, random_file_name)

        image_path = image.path  # 将路径设置为为实际图像文件的路径
        image = load_image(image_path)
        # files = {'file': image}
        logger.info(image)
        logger.info("6" * 10)

        remote_flask_url = 'http://125.220.157.228:29990/inpaint/upload'
        response_image_saved_path_and_name = upload_image_to_remote_flask(image_path, remote_flask_url)
        # response_image_saved_path_and_name = requests.post(remote_flask_url, files=files)
        logger.info("7" * 10)
        logger.info("response_image_saved_path_and_name" + response_image_saved_path_and_name)
        response_data = json.loads(response_image_saved_path_and_name)
        file_path = response_data["message"]
        logger.info("file_path" + "---" + file_path)
        logger.info("8" * 10)
        data = {'input_img': file_path,
                'coords_type': "key_in",
                'point_coords': args['point_coords'],
                'point_labels': args['point_labels'],
                'dilate_kernel_size': 10,  # fixme
                'output_dir': './results'
                }
        logger.info("point_coords" + str(data['point_coords']))

        remote_flask_remove_replace_anything_url = 'http://125.220.157.228:29990/inpaint/remove'
        response_remove_thing = requests.post(remote_flask_remove_replace_anything_url, json=data)

        logger.info("response_fill_thing: type: %s", type(response_remove_thing))

        response_remove_thing = response_remove_thing.json()
        response_remove_thing_image_bytes = response_remove_thing['image']

        image_data = base64.b64decode(response_remove_thing_image_bytes)
        responsed_image = Image.open(io.BytesIO(image_data))
        logger.info("13" * 10)
        shutil.copyfile(image_path, remove_filepath)
        responsed_image.save(remove_filepath)
        image = ImageModel.create_from_path(remove_filepath)
        image.enhanced = True
        image.deleted = False
        image.save()
        return query_util.fix_ids(image)


@api.route('/rename')
class ImageRename(Resource):

    @api.expect(image_rename_args)
    def post(self):
        args = image_rename_args.parse_args()
        image_id = args['image_id']
        new_filename = args['new_filename']

        image = current_user.images.filter(id=image_id, deleted=False).first()
        if image:
            try:
                # 获取图片文件的扩展名
                _, ext = os.path.splitext(image.file_name)
                # 构造新的文件名
                new_file_name = f"{new_filename}{ext}"
                # 获取文件的目录路径
                file_directory = os.path.dirname(image.path)
                # 构造新的文件路径
                new_file_path = os.path.join(file_directory, new_file_name)

                # 重命名文件
                os.rename(image.path, new_file_path)

                # 更新数据库中的文件名字段
                image.file_name = new_file_name
                image.path = new_file_path
                image.save()

                return {'message': 'Image file renamed successfully'}, 200
            except Exception as e:
                return {'message': f'Failed to rename image file: {str(e)}'}, 500
        else:
            return {'message': 'Image not found or already deleted'}, 404


@api.route('/delete')
class ImageDeleted(Resource):

    @api.expect(image_deleted_by_id)
    def post(self):
        args = image_deleted_by_id.parse_args()
        image_id = args['image_id']
        image = current_user.images.filter(id=image_id, deleted=False).first()
        image.deleted = True
        image.save()

        return "delete image successed"

@api.route('/remove_by_segmentation')
class ImageRemovedBySegmentation(Resource):

    @api.expect(image_remove_by_segmentatino)
    def post(self):
        args = image_remove_by_segmentatino.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        annotation_id = args['annotation_id']
        annotation = AnnotationModel.objects(id=annotation_id).first()
        #获取segmentation
        #fixme:这地方需要修改

        # segmentation = annotation.segmentation[3]

        # 找到最长的 segmentation
        longest_segmentation = max(annotation.segmentation, key=len)
        # 获取最长的 segmentation
        segmentation = longest_segmentation

        logger.info("="*30)
        logger.info("segmentation:" + str(segmentation))
        logger.info("=" * 30)
        if not image:
            return {"message": "Invalid image ID"}, 400
        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory
        remove_directory = dataset_path + r'removed/'
        if not os.path.exists(remove_directory):
            os.makedirs(remove_directory)
        random_file_name = str(uuid.uuid4()) + ".jpg"
        remove_filepath = os.path.join(remove_directory, random_file_name)
        image_path = image.path  # 将路径设置为为实际图像文件的路径
        image = load_image(image_path)
        # files = {'file': image}
        logger.info(image)
        logger.info("6" * 10)

        remote_flask_url = 'http://125.220.157.228:29990/inpaint/upload'
        response_image_saved_path_and_name = upload_image_to_remote_flask(image_path, remote_flask_url)
        # response_image_saved_path_and_name = requests.post(remote_flask_url, files=files)

        logger.info("7" * 10)
        logger.info("response_image_saved_path_and_name" + response_image_saved_path_and_name)
        response_data = json.loads(response_image_saved_path_and_name)
        file_path = response_data["message"]
        logger.info("file_path" + "---" + file_path)
        logger.info("8" * 10)
        data = {'input_img': file_path,
                'segmentation': segmentation,
                'dilate_kernel_size': 10,
                'output_dir': './results'
                }

        remote_flask_remove_replace_anything_url = 'http://125.220.157.228:29990/inpaint/remove_by_segmentation'
        response_remove_thing = requests.post(remote_flask_remove_replace_anything_url, json=data)

        logger.info("response_fill_thing: type: %s", type(response_remove_thing))

        response_remove_thing = response_remove_thing.json()
        response_remove_thing_image_bytes = response_remove_thing['image']

        image_data = base64.b64decode(response_remove_thing_image_bytes)
        responsed_image = Image.open(io.BytesIO(image_data))
        logger.info("13" * 10)
        shutil.copyfile(image_path, remove_filepath)
        responsed_image.save(remove_filepath)
        image_removed = ImageModel.create_from_path(remove_filepath)
        image_removed.enhanced = True
        image_removed.deleted = False

        image_removed.save()
        return query_util.fix_ids(image_removed)


@api.route('/fill_by_segmentation')
class ImageFilledBySegmentation(Resource):
    @api.expect(image_fill_by_segmentatino)
    def post(self):
        args = image_fill_by_segmentatino.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        annotation_id = args['annotation_id']
        text_prompt = args['text_prompt']
        annotation = AnnotationModel.objects(id=annotation_id).first()

        # 找到最长的 segmentation
        longest_segmentation = max(annotation.segmentation, key=len)
        # 获取最长的 segmentation
        segmentation = longest_segmentation

        logger.info("="*30)
        logger.info("segmentation:" + str(segmentation))
        logger.info("=" * 30)
        if not image:
            return {"message": "Invalid image ID"}, 400
        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory
        # fill_filepath
        fill_directory = dataset_path + r'filled/'
        if not os.path.exists(fill_directory):
            os.makedirs(fill_directory)
        random_file_name = str(uuid.uuid4()) + ".jpg"
        remove_filepath = os.path.join(fill_directory, random_file_name)

        image_path = image.path  # 将路径设置为为实际图像文件的路径
        image = load_image(image_path)

        logger.info(image)
        logger.info("6" * 10)

        remote_flask_url = 'http://125.220.157.228:29990/inpaint/upload'
        response_image_saved_path_and_name = upload_image_to_remote_flask(image_path, remote_flask_url)

        logger.info("response_image_saved_path_and_name" + response_image_saved_path_and_name)
        response_data = json.loads(response_image_saved_path_and_name)
        file_path = response_data["message"]
        logger.info("file_path" + "---" + file_path)
        logger.info("8" * 10)
        data = {'input_img': file_path,
                'segmentation': segmentation,
                'text_prompt': text_prompt,
                'dilate_kernel_size': 10,
                'output_dir': './results'
                }

        remote_flask_remove_replace_anything_url = 'http://125.220.157.228:29990/inpaint/fill_anything_with_segmentation'
        response_remove_thing = requests.post(remote_flask_remove_replace_anything_url, json=data)

        logger.info("response_fill_thing: type: %s", type(response_remove_thing))

        response_remove_thing = response_remove_thing.json()
        response_remove_thing_image_bytes = response_remove_thing['image']

        image_data = base64.b64decode(response_remove_thing_image_bytes)
        responsed_image = Image.open(io.BytesIO(image_data))
        logger.info("13" * 10)
        shutil.copyfile(image_path, remove_filepath)
        responsed_image.save(remove_filepath)
        image = ImageModel.create_from_path(remove_filepath)
        image.enhanced = True
        image.deleted = False
        image.save()
        return query_util.fix_ids(image)


# 嵌入水印
@api.route('/embed_watermark')
class ImageEmbedWatermark(Resource):
    # embed_watermark.add_argument('input_img_id', type=int, required=True)
    # embed_watermark.add_argument('output_img_path', type=str, required=True)
    @api.expect(embed_watermark)
    def post(self):
        args = embed_watermark.parse_args()
        input_img_id = args['input_img_id']
        # output_img_path = args['output_img_path']

        # image = ImageModel.objects(id=image_id).first()
        image = current_user.images.filter(id=input_img_id, deleted=False).first()
        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory

        watermark_directory = dataset_path + r'watermarked/'
        if not os.path.exists(watermark_directory):
            os.makedirs(watermark_directory)
        random_file_name = str(uuid.uuid4()) + ".jpg"
        watermark_filepath = os.path.join(watermark_directory, random_file_name)

        with open(image.path, 'rb') as image_file:
            files = {'image': image_file}
            remote_flask_url = 'http://125.220.157.228:29990/embed_watermark'
            data = {
                'output_img_path': watermark_filepath
            }
            r = requests.post(remote_flask_url, data=data, files=files)

        return "successed"


@api.route('/extract_watermark')
class ImageExtractWatermark(Resource):
#     extract_watermark = reqparse.RequestParser()
#     extract_watermark.add_argument('input_img_with_watermark_id', type=int, required=True)
    @api.expect(extract_watermark)
    def post(self):
        args = extract_watermark.parse_args()
        input_img_id = args['input_img_with_watermark_id']

        image = current_user.images.filter(id=input_img_id, deleted=False).first()

        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory

        with open(image.path, 'rb') as image_file:
            files = {'image': image_file}
            remote_flask_url = 'http://125.220.157.228:29990/extract_watermark'
            r = requests.post(remote_flask_url,files=files)

        response_data = r.json()
        watermark = response_data.get('watermark', '')
        logger.info(watermark)
        return watermark



@api.route('/enhance_by_instruct_pix_to_pix')
class ImageRemovedByPix2Pix(Resource):
    # pix_to_pix = reqparse.RequestParser()
    # pix_to_pix.add_argument('image_id', type=int, required=True)
    # pix_to_pix.add_argument('text_prompt', type=str, default=None, required=True)
    @api.expect(pix_to_pix)
    def post(self):
        args = pix_to_pix.parse_args()
        image_id = args['image_id']
        image = ImageModel.objects(id=image_id).first()
        text_prompy = args['text_prompt']

        dataset = image.dataset_id
        dataset_path = DatasetModel.objects(id=dataset).first().directory
        enhanced_directory = dataset_path + r'enhanced/'
        if not os.path.exists(enhanced_directory):
            os.makedirs(enhanced_directory)
        random_file_name = str(uuid.uuid4()) + ".jpg"
        enhanced_filepath = os.path.join(enhanced_directory, random_file_name)
        image_path = image.path  # 将路径设置为为实际图像文件的路径
        image = load_image(image_path)
        with open(image_path, 'rb') as image_file:
            logger.info("type image_file:" + str(type(image_file)))
            files = {'file': image_file}
            remote_flask_url = 'http://125.220.157.228:29994/edit_image'
            data = {
                'edit': text_prompy
            }
            r = requests.post(remote_flask_url, data=data, files=files)

        response_data = r.json()
        edited_image_base64 = response_data["edited_image"]
        # 将 Base64 编码的图像数据解码为字节数据
        image_data = base64.b64decode(edited_image_base64)

        responsed_image = Image.open(io.BytesIO(image_data))
        shutil.copyfile(image_path, enhanced_filepath)
        responsed_image.save(enhanced_filepath)
        image = ImageModel.create_from_path(enhanced_filepath)
        image.enhanced = True
        image.deleted = False
        image.save()
        return query_util.fix_ids(image)




# @api.route('/auto')
# class AnnotationAuto(Resource):
#
#     @api.expect(create_auto)
#     def post(self):
#         args = create_auto.parse_args()
#         image_id = args.get('image_id')
#         category_id = args.get('category_id')
#         isbbox = args.get('isbbox')
#         metadata = args.get('metadata', {})
#         keypoints = args.get('keypoints', [])
#         x = args.get('x')
#         y = args.get('y')
#
#         image = ImageModel.objects(id=image_id).first()
#         with open(image.path, 'rb') as image_file:
#             files = {'image': image_file}
#             remote_flask_url = 'http://125.220.157.228:29991/single'
#             data = {
#                 'x': x,
#                 'y': y
#             }
#             r = requests.post(remote_flask_url, data=data, files=files)
#         logger.info(r.json())
#         response = r.json()
#         annotations = response['annotations']
#         new_annotations = []
#         for annotation in annotations:
#             segmentation = annotation['segmentation']
#             logger.info(segmentation)
#             area = annotation['area']
#             bbox = annotation['bbox']
#             try:
#                 new_annotation = AnnotationModel(
#                     image_id=image_id,
#                     category_id=category_id,
#                     metadata=metadata,
#                     segmentation=segmentation,
#                     bbox=bbox,
#                     area=area,
#                     keypoints=keypoints,
#                     isbbox=isbbox
#                 )
#                 new_annotation.save()
#                 new_annotations.append(query_util.fix_ids(new_annotation))
#             except (ValueError, TypeError) as e:
#                 return {'message': str(e)}, 400
#         res = {
#             "annotations": new_annotations
#         }
#         # return json.dumps(res)
#         return res, 200

# 提取水印
# @api.route('/extract_watermark')
# class ImageExtractWatermark(Resource):
#     extract_watermark = reqparse.RequestParser()
#     extract_watermark.add_argument('input_img_with_watermark_id', type=int, required=True)
#     @api.expect(extract_watermark)
#     def post(self):
#         args = extract_watermark.parse_args()
#         input_img_with_watermark_id = args['input_img_with_watermark_id']
#
#         image = current_user.images.filter(id=input_img_with_watermark_id, deleted=False).first()
#         result = extract_watermark_str(image.path)
#
#         return result